# ''' Python 代码开始

import cv2
import numpy as np
import glob
import yaml
import os
import argparse
import abc

# ===================================================================
# 策略模式: 定义标定板处理器的基类和具体实现
# (这部分与您的CameraCalibrator脚本中的完全相同)
# ===================================================================

class CalibrationPattern(abc.ABC):
    """标定板处理器的抽象基类 (接口)."""
    def __init__(self, params):
        self.params = params
        self._validate_params()

    @abc.abstractmethod
    def _validate_params(self):
        pass

    @abc.abstractmethod
    def get_object_points(self):
        pass

    @abc.abstractmethod
    def find_corners(self, image):
        pass

    def draw_corners(self, image, corners, ret):
        """在图像上绘制找到的角点 (通用实现)."""
        return cv2.drawChessboardCorners(image, self.pattern_size, corners, ret)

class CheckerboardPattern(CalibrationPattern):
    """棋盘格标定板处理器."""
    def _validate_params(self):
        if not all(k in self.params for k in ['cols', 'rows', 'square_size']):
            raise ValueError("棋盘格模式需要 'cols', 'rows', 'square_size' 参数.")
        self.pattern_size = (self.params['cols'], self.params['rows'])
        self.square_size = self.params['square_size']

    def get_object_points(self):
        objp = np.zeros((self.pattern_size[0] * self.pattern_size[1], 3), np.float32)
        objp[:, :2] = np.mgrid[0:self.pattern_size[0], 0:self.pattern_size[1]].T.reshape(-1, 2)
        return objp * self.square_size

    def find_corners(self, image):
        flags = cv2.CALIB_CB_ADAPTIVE_THRESH + cv2.CALIB_CB_NORMALIZE_IMAGE
        ret, corners = cv2.findChessboardCorners(image, self.pattern_size, flags=flags)
        if not ret:
             ret, corners = cv2.findChessboardCorners(image, self.pattern_size, cv2.CALIB_CB_FAST_CHECK)
        return ret, corners

class CirclesGridPattern(CalibrationPattern):
    """对称圆形网格标定板处理器."""
    def _validate_params(self):
        if not all(k in self.params for k in ['cols', 'rows', 'spacing']):
            raise ValueError("对称圆形网格模式需要 'cols', 'rows', 'spacing' 参数.")
        self.pattern_size = (self.params['cols'], self.params['rows'])
        self.spacing = self.params['spacing']

    def get_object_points(self):
        objp = np.zeros((self.pattern_size[0] * self.pattern_size[1], 3), np.float32)
        objp[:, :2] = np.mgrid[0:self.pattern_size[0], 0:self.pattern_size[1]].T.reshape(-1, 2)
        return objp * self.spacing

    def find_corners(self, image):
        flags = cv2.CALIB_CB_SYMMETRIC_GRID
        return cv2.findCirclesGrid(image, self.pattern_size, flags=flags)

# ===================================================================
# 主位姿估计类: PoseEstimator
# ===================================================================

class PoseEstimator:
    """
    使用已知的相机内参来计算标定板位姿的类.

    核心功能:
    - 通过YAML配置文件加载所有参数。
    - 加载由相机标定生成的内参和畸变系数。
    - 采用策略模式支持多种标定板。
    - 处理单个或多个图像。
    - 对每张图像使用 cv2.solvePnP 计算位姿 (旋转和平移向量)。
    - 将结果保存为4x4齐次变换矩阵，可直接用于手眼标定。
    - 生成带有3D坐标轴的可视化图像以供验证。
    """
    
    PATTERN_HANDLERS = {
        'checkerboard': CheckerboardPattern,
        'circles': CirclesGridPattern,
        # 可以添加 'asymmetric_circles' 等其他模式
    }

    def __init__(self, config_path):
        """
        初始化, 加载配置、相机内参和标定板处理器.
        :param config_path: 配置文件 (pose_config.yaml) 的路径.
        """
        print("INFO: 初始化位姿估计器...")
        with open(config_path, 'r') as f:
            self.config = yaml.safe_load(f)
        
        # 加载配置
        self.intrinsics_path = self.config['camera_intrinsics_path']
        self.image_path_pattern = self.config['image_path']
        self.output_path = self.config['output_path']
        self.save_verification = self.config.get('save_verification_image', True)
        
        # 加载相机内参
        self._load_intrinsics()

        # 初始化标定板处理器
        self.pattern_type = self.config['pattern_type']
        self.pattern_params = self.config['pattern_params']
        handler_class = self.PATTERN_HANDLERS.get(self.pattern_type)
        if not handler_class:
            raise ValueError(f"错误: 不支持的标定板类型 '{self.pattern_type}'.")
        self.pattern_handler = handler_class(self.pattern_params)
        print(f"INFO: 已选择标定板类型: '{self.pattern_type}'")

        # 创建输出目录
        if not os.path.exists(self.output_path):
            os.makedirs(self.output_path)
            print(f"INFO: 已创建输出目录: {self.output_path}")
        
        if self.save_verification:
            self.verification_path = os.path.join(self.output_path, 'verification_outputs')
            if not os.path.exists(self.verification_path):
                os.makedirs(self.verification_path)
            print(f"INFO: 验证图像将保存至: {self.verification_path}")

    def _load_intrinsics(self):
        """从 .npz 或 .yaml 文件加载相机内参."""
        print(f"INFO: 正在从 {self.intrinsics_path} 加载相机内参...")
        if self.intrinsics_path.endswith('.npz'):
            data = np.load(self.intrinsics_path)
            self.mtx = data['mtx']
            self.dist = data['dist']
        elif self.intrinsics_path.endswith('.yaml'):
            with open(self.intrinsics_path, 'r') as f:
                data = yaml.safe_load(f)
            self.mtx = np.array(data['camera_matrix'])
            self.dist = np.array(data['distortion_coefficients'])
        else:
            raise ValueError("错误: 内参文件格式必须是 .npz 或 .yaml")
        print("  - 内参矩阵 (mtx):\n", self.mtx)
        print("  - 畸变系数 (dist):\n", self.dist)

    def estimate_pose(self):
        """
        执行位姿估计的主流程.
        """
        print("\nINFO: 开始位姿估计流程...")
        images = sorted(glob.glob(self.image_path_pattern))
        if not images:
            print(f"ERROR: 在路径 '{self.image_path_pattern}' 未找到任何图像. 请检查配置文件.")
            return

        print(f"INFO: 找到 {len(images)} 张图像进行处理.")
        
        objp = self.pattern_handler.get_object_points()

        for fname in images:
            img = cv2.imread(fname)
            if img is None:
                print(f"WARNING: 无法读取图像 {fname}, 已跳过.")
                continue

            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            base_filename = os.path.basename(fname)
            print(f"--- 正在处理图像: {base_filename} ---")
            
            ret, corners = self.pattern_handler.find_corners(gray)
            
            if ret:
                print("  - 成功找到角点.")
                # 使用 solvePnP 计算位姿
                # 它计算的是从标定板坐标系到相机坐标系的变换
                ret_pnp, rvec, tvec = cv2.solvePnP(objp, corners, self.mtx, self.dist)
                
                if ret_pnp:
                    print("  - 成功通过 solvePnP 计算位姿.")
                    self._save_results(rvec, tvec, base_filename)
                    if self.save_verification:
                        self._generate_verification_image(img, corners, rvec, tvec, base_filename)
                else:
                    print(f"  - ERROR: cv2.solvePnP 未能成功计算位姿.")
            else:
                print(f"  - ERROR: 未能找到角点, 跳过此图像.")
        
        print("\n位姿估计流程完成！")

    def _save_results(self, rvec, tvec, base_filename):
        """将位姿结果保存为4x4矩阵文件."""
        # 将旋转向量转换为旋转矩阵
        R, _ = cv2.Rodrigues(rvec)
        
        # 创建4x4齐次变换矩阵 (target_H_cam)
        target_H_cam = np.identity(4)
        target_H_cam[:3, :3] = R
        target_H_cam[:3, 3] = tvec.flatten()

        print("  - 计算出的位姿 (target_H_cam):")
        np.set_printoptions(precision=4, suppress=True)
        print(target_H_cam)

        # 保存为 .txt 文件, 可直接用于手眼标定脚本
        output_filename = os.path.splitext(base_filename)[0] + '.txt'
        save_path = os.path.join(self.output_path, output_filename)
        np.savetxt(save_path, target_H_cam, fmt='%.8f')
        print(f"  - 位姿已保存至: {save_path}")

    def _generate_verification_image(self, image, corners, rvec, tvec, base_filename):
        """生成并保存带有坐标轴的验证图像."""
        axis_length = 3 * self.pattern_handler.params.get('square_size', self.pattern_handler.params.get('spacing', 1.0))
        
        # 在图像上绘制3D坐标轴
        # cv2.drawFrameAxes 在OpenCV 4.7.0+ 中可用
        # 对于旧版本，可使用 cv2.aruco.drawAxis
        try:
            drawn_img = cv2.drawFrameAxes(image, self.mtx, self.dist, rvec, tvec, axis_length)
        except AttributeError:
            print("WARNING: cv2.drawFrameAxes 不可用. 请考虑升级OpenCV. 跳过坐标轴绘制。")
            drawn_img = self.pattern_handler.draw_corners(image, corners, True)

        save_path = os.path.join(self.verification_path, f'pose_verification_{base_filename}')
        cv2.imwrite(save_path, drawn_img)
        print(f"  - 验证图像已保存至: {save_path}")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="使用已知相机内参计算标定板位姿")
    parser.add_argument('--config', type=str, default='pose_config.yaml', 
                        help='配置文件的路径 (默认为: pose_config.yaml)')
    args = parser.parse_args()
    
    if not os.path.exists(args.config):
        print(f"ERROR: 配置文件 '{args.config}' 不存在. 请确保该文件存在或通过 --config 参数指定正确路径。")
    else:
        try:
            estimator = PoseEstimator(args.config)
            estimator.estimate_pose()
        except (ValueError, FileNotFoundError, KeyError) as e:
            print(f"\nERROR: 处理失败. 可能是配置文件内容有误或文件路径问题。")
            print(f"  > 错误详情: {e}")

# ''' Python代码结束